# -*- coding: utf-8 -*-
"""
@author:xuyuntao
@time:2021/1/17:8:36
@email:xuyuntao@189.cn
"""
import numpy
import numpy.fft as nfft
import scipy.signal as scisignal
import types

class DigitalModulationGenerator(object):
    """数字调制生成器：
    可以生成各种简单的数字调制信号，并附有滤波的函数"""
    def __init__(self,bitsData:numpy.ndarray,bytesRate:int=1,carrierFreq:int=1000,sampleFreq:int=5000):
        """初始化函数
        bitsData：比特数据，按进制不同转换为不同位码元
        bytesRate：码元速率，单位Hz
        carrierFreq：载波频率，单位Hz
        sampleFreq：采样频率，单位Hz"""
        if type(bitsData)==numpy.ndarray:   # 还需要检验比特数据是否只有1、0
            self.bitsData = bitsData.astype(numpy.bool).astype(numpy.int8).reshape(-1)  # 比特数据，按进制不同转换为不同位码元
        else:
            if bitsData==None:
                self.bitsData=numpy.random.randint(0, 2, 1024,dtype=numpy.int8) # 没有输入情况下生成1024个比特
            else:
                raise ValueError("输入比特数据类型错误")
        self.bytesRate=round(bytesRate) # 码元速率
        self.carrierFreq=round(carrierFreq) # 载波频率
        if (sampleFreq<2*(carrierFreq+bytesRate)): # 低通采样最低频率
            raise ValueError("采样频率低于信号最高频率的两倍")
        self.sampleFreq=round(sampleFreq) # 采样频率
        if (self.sampleFreq/self.bytesRate)%1>0:  # 不能实现一个码元采样数为整数，不可实现
            raise ValueError("一个码元周期内的采样点数不为整数")
        self.NperBytes=round(self.sampleFreq / self.bytesRate) # 一个码元对应的点数
        self.timeLength = self.bitsData.size / self.bytesRate # 时间长度，单位：秒
        self.timeN=self.bitsData.size*self.NperBytes # 时间序列长度，单位：个
        self.bytesPeriod = 1 / self.bytesRate  # 基带信号周期
        self.samplePeriod = 1 / self.sampleFreq  # 采样间隔
        self.seperateFreq=None    # FSK调制的频率间隔

        self.timeSeries = numpy.linspace(0, self.timeLength, self.timeN+1)  # 时间序列，长度为比特数/码元速率
        # self.timeSeries = numpy.arange(0,self.timeLength,self.samplePeriod)
        self.carrier,self.carrier_Complex=None,None  # 初始化载波缓冲
        self._carrierGet()  # 计算载波，后期考虑加入相位噪声生成

        self.signal=None # 调制信号
        self.Constel=None # 基带经频谱搬移后的信号
        self.signalN=None # 信号点数

        self.signalFilter=None # 滤波后信号

    def _getBytesNum(self,bitsData,binLen=1):
        """将输入比特序列转化成以十进制数表示的M进制数组
        bitsData:比特数据
        binLen:多进制数的二进制表示的二进制长度"""
        bytes = bitsData.reshape([-1, binLen])  # 转换为多位二进制表示
        power=numpy.arange(0,binLen,dtype=numpy.int16).reshape([1,-1])\
            .repeat(bytes.shape[0],axis=0)  # 缓存次方矩阵
        bytesData=bytes*(2**power)  # 表示以二进制表示的M进制数
        bytesNum=bytesData.sum(axis=1)  # 换算成十进制数表示
        return bytesNum

    def _dataLengthValid(self,M=8):
        """输入数据长度检验，需满足二进制转M进制的位数要求
        返回值为log2（M），即M需要多少位二进制可表示"""
        if M%1>0:   # 非整数
            raise ValueError("M不为整数")
        if M<=0:    # 非正数
            raise ValueError("M不为正数")
        if M>=32767: # 太大了
            raise ValueError("M大于2^16")
        if numpy.log2(M)%1>0:  # 不是2的次方
            raise ValueError("log2(M)不为整数")
        if (self.bitsData.size/numpy.log2(M))%1>0:  # 比特数不满足多进制转换要求
            raise ValueError("输入比特数不满足转换M进制要求")
        return round(numpy.log2(M))

    def _carrierGet(self,lineWidth=0):
        """生成载波，
        后期考虑根据lineWidth线宽生成带相位噪声的载波"""
        # todo:后期加入相位噪声生成
        self.carrier_Complex = numpy.exp(numpy.complex(0, 1) * \
                                         (2 * numpy.pi * self.timeSeries * self.carrierFreq))
        self.carrier = self.carrier_Complex.real#numpy.cos(2 * numpy.pi * self.carrierFreq * self.timeSeries)
        return self.carrier,self.carrier_Complex
#------------------psk------------------
    def _pskConstel(self,M,bytesNum,initPhase=0):
        """生成PSK的基带复信号
        M:阶数
        bytesNum:码元数据，取值应为0~(M-1)
        initPhase:初始相位角"""
        _theta = 2 * numpy.pi * bytesNum / M # 根据调制阶数和码元数据求相位角
        pskConstel = numpy.exp(numpy.complex(0,1) * (_theta + initPhase))  # exp(j*theta)
        return pskConstel

# ------------------qam------------------
    def _qamConstel(self,binLen,bitsData):
        """生成QAM的基带复信号
        需注意，M=2^2k, k=1,2,3...
        binLen:多进制数的二进制表示的二进制长度，即2k
        bitsData:比特数据，取值为0~1"""
        if binLen<=0:   # 指数为负
            raise ValueError("M数据错误")
        if binLen/2%1>0:  # 指数非2的倍数，不满足qam条件
            raise ValueError("M不等于二的双数指数次，即要求M=2^2k, (k=1,2,3...)")
        halfBinLen=round(binLen/2)
        vect = numpy.arange(1, 2**halfBinLen, 2)
        symbols = numpy.concatenate((numpy.flip(-1*vect, axis=0), vect)).astype(int)
        byteData_tmp=self._getBytesNum(bitsData,halfBinLen).reshape([-1,2])
        bytesData_I=byteData_tmp[:,0].reshape(-1).astype(numpy.int16)
        bytesData_Q=byteData_tmp[:,1].reshape(-1).astype(numpy.int16)
        constel_I = symbols[bytesData_I]
        constel_Q = symbols[bytesData_Q]
        qamConstel=constel_I+numpy.complex(0,1)*constel_Q
        return qamConstel

# ------------------fsk------------------
    def _fskConstel(self,M,bytesNum,seperateFreq,NperBytes,sampleFreq,phaseContinue=True):
        """生成FSK的基带复信号
        M:阶数
        bytesNum:码元数据，取值应为0~(M-1)
        seperateFreq:每个数的频点之间的频差
        NperBytes:每个码元周期内的采样点数
        sampleFreq:采样速率
        phaseContinue:是否为连续相位，相位连续需要seperateFreq/(2*symbolRate)为整数
        """
        maxFreq = ((M - 1) / 2) * seperateFreq
        if (maxFreq > sampleFreq / 2):  # 检验采样频率
            raise ValueError("采样频率低于fsk基带信号最高频率的两倍")

        phaseIncr=numpy.matmul(numpy.arange(NperBytes).reshape([-1, 1]), \
                               numpy.arange(-1 * (M - 1), M, 2).reshape([1, -1])) \
                  * 2 * numpy.pi * seperateFreq / 2 / sampleFreq
        phIncrSym = phaseIncr[-1,:]
        phIncrSamp = phaseIncr[1,:]
        OscPhase = numpy.zeros(M)

        Phase = numpy.zeros([bytesNum.size,NperBytes])

        if phaseContinue or ((NperBytes/sampleFreq/2*seperateFreq)%1>0):
            prevPhase = 0
            for iSym in range(bytesNum.size):
                if phaseContinue:
                    ph1 = prevPhase
                else:
                    ph1 = OscPhase[bytesNum[iSym]]
                Phase[iSym,:]=ph1+phaseIncr[:,bytesNum[iSym]]
                if phaseContinue:
                    pass
                else:
                    OscPhase=numpy.remainder(OscPhase+phIncrSym+phIncrSamp,2*numpy.pi)
                prevPhase=Phase[iSym,-1]+phIncrSamp[bytesNum[iSym]]

            fskConstel=numpy.exp(numpy.complex(0,1) * Phase).reshape(-1)
        else:
            exp_phaseIncr = numpy.exp(numpy.complex(0,1) * phaseIncr)
            fskConstel=exp_phaseIncr[:,bytesNum].reshape(-1)

        return fskConstel

# ------------------ask------------------
    def _askConstel_unipolar(self, bytesNum):
        """bytesData:输入码元数据，取值应为0-阶数，此函数生成单极性ask基带"""
        return bytesNum

    def _askConstel_bipolar(self,bytesNum,M,dataType=numpy.int32):
        """bytesData:输入码元数据，取值应为0-阶数，此函数生成双极性ask基带"""
        constel_tmp = numpy.zeros(M, dtype=dataType)
        half_tmp = numpy.arange(1 - M, 0, 2)
        constel_tmp[:round(M / 2)] = half_tmp
        constel_tmp[round(M / 2):] = -1 * half_tmp[::-1]
        constel = constel_tmp[bytesNum]
        return constel

# ------------------统一输入调制方式------------------
    def modulate_ask(self,binLen,M,polar=True):
        """ASK星座图生成的通用函数
        polar=True表示双极性，反之表示单极性"""
        bytesNum = self._getBytesNum(self.bitsData,binLen)
        if polar:
            constel=self._askConstel_bipolar(bytesNum,M,numpy.int32)
        else:
            constel=self._askConstel_unipolar(bytesNum)
        return constel.repeat(self.NperBytes)

    def modulate_psk(self,binLen,M, initPhase=0):
        """PSK星座图生成的通用函数"""
        bytesNum = self._getBytesNum(self.bitsData,binLen)
        constel=self._pskConstel(M, bytesNum.repeat(self.NperBytes), initPhase)
        return constel

    def modulate_fsk(self,binLen,M,seperateFreq,phaseContinue=True):
        """FSK星座图生成的通用函数"""
        bytesNum = self._getBytesNum(self.bitsData, binLen)
        constel=self._fskConstel(M,bytesNum,seperateFreq,self.NperBytes,self.sampleFreq,phaseContinue)
        return constel

    def modulate_qam(self,binLen,M):
        """QAM星座图生成的通用函数"""
        constel=self._qamConstel(binLen,self.bitsData).repeat(self.NperBytes)
        return constel

# ------------------为方便统一调用，设置fskSep设置函数------------------
    def setFSKSeperateFreq(self,seperateFreq):
        """FSK频率间隔提前设置函数，为方便统一使用getSignal生成调制信号，
        提供直接设置FSK频率间隔的函数，如果getSignal处没有输入seperateFreq参数，
        就调用这个函数提前设置的值，如果输入了就采用输入值，
        如果都没有，则报错
        seperateFreq: FSK调制频率间隔"""
        self.seperateFreq=seperateFreq
        return self.seperateFreq

# ------------------信号功率归一化------------------
    def _normalize(self,signal):
        """信号幅度归一化函数"""
        return signal / numpy.sqrt(numpy.mean(numpy.abs(signal) ** 2))

# ------------------获取指定的调制方式------------------
    def getConstel(self,modulate_method,M,norm=True,*args,**kwargs):
        """获得指定调制信号的复基带信号
        modulate_method:调制方式，函数类型，如：DMG.modulate_psk
        M:阶数，必须为整数，且为2的指数次
        还有一些个别调制方法的特殊参数在最后输入
        输出为复基带信号"""
        M=round(M)
        self.signalFilter=None
        binLen=self._dataLengthValid(M)
        self.signalN = round(self.bitsData.size/binLen)*self.NperBytes
        if type(modulate_method)==types.MethodType:
            Constel=modulate_method(binLen,M,*args,**kwargs)
        elif type(modulate_method)==str:
            if modulate_method.upper()=="ASK":
                Constel=self.modulate_ask(binLen,M,*args,**kwargs)
            elif modulate_method.upper()=="PSK":
                Constel=self.modulate_psk(binLen,M,*args,**kwargs)
            elif modulate_method.upper()=="FSK":
                if (len(args)>=1) or ("seperateFreq" in kwargs.keys()):
                    Constel=self.modulate_fsk(binLen,M,*args,**kwargs)
                else:
                    if (self.seperateFreq==None):
                        raise ValueError("没有设置seperateFreq参数")
                    else:
                        Constel = self.modulate_fsk(binLen, M,self.seperateFreq)
            elif modulate_method.upper()=="QAM":
                Constel=self.modulate_qam(binLen,M)
            else:
                raise ValueError("输入类型为DMG类不支持的调制类型")
        else:
            raise TypeError("modulate_method输入参数类型必须为DMG的内建方法或者为str类型")

        self.Constel=Constel

        if norm:
            return self._normalize(Constel.copy())
        else:
            return Constel.copy()

    def _getSignal(self,modulate_method,M,*args,**kwargs):
        """利用星座图信号与载波相乘获得指定调制信号
        modulate_method:调制方式，函数类型，如：DMG.modulate_psk
        或者为str类型，可选：ASK、FSK、PSK、QAM
        M:阶数，必须为整数，且为2的指数次
        还有一些个别调制方法的特殊参数在最后输入
        输出为调制后的复信号"""
        self.getConstel(modulate_method,M,False,*args,**kwargs)

        self.signal=self.carrier_Complex[:self.signalN].conjugate()*self.Constel
        # 基带调制到信号，采用基带复信号与复载波的共轭相乘
        return self.signal

    def getSignal_real(self, modulate_method, M, norm=True, *args, **kwargs):
        """利用星座图信号与载波相乘获得指定调制信号
        modulate_method:调制方式，函数类型，如：DMG.modulate_psk
        M:阶数，必须为整数，且为2的指数次
        还有一些个别调制方法的特殊参数在最后输入
        输出为调制后的实信号和时间序列"""
        self._getSignal(modulate_method,M,*args,**kwargs)
        if norm:
            return self._normalize(self.signal.real.copy()),self.timeSeries[:self.signalN].copy()
        else:
            return self.signal.real.copy(), self.timeSeries[:self.signalN].copy()

    def getSignal(self, modulate_method, M, norm=True, *args, **kwargs):
        """利用星座图信号与载波相乘获得指定调制信号
        modulate_method:调制方式，函数类型，如：DMG.modulate_psk
        M:阶数，必须为整数，且为2的指数次
        还有一些个别调制方法的特殊参数在最后输入
        输出为调制后的复信号和时间序列"""
        self._getSignal(modulate_method,M,*args,**kwargs)
        if norm:
            return self._normalize(self.signal.copy()),self.timeSeries[:self.signalN].copy()
        else:
            return self.signal.copy(), self.timeSeries[:self.signalN].copy()

    def get_Constel_Signal(self, modulate_method, M, norm=True, *args, **kwargs):
        """利用星座图信号与载波相乘获得指定调制信号
        modulate_method:调制方式，函数类型，如：DMG.modulate_psk
        M:阶数，必须为整数，且为2的指数次
        还有一些个别调制方法的特殊参数在最后输入
        输出为调制后的复信号和时间序列"""
        self._getSignal(modulate_method,M,*args,**kwargs)
        if norm:
            return self._normalize(self.Constel.copy()),self._normalize(self.signal.copy()),self.timeSeries[:self.signalN].copy()
        else:
            return self.Constel.copy(),self.signal.copy(), self.timeSeries[:self.signalN].copy()

# ------------------滤波器------------------
    def filter(self,halfBand,transitionBand,passAmp=2,stopAmp=23):
        """内置的简单的滤波器函数
        滤波器的中心频率自动设为载波频率
        halfBand：滤波带宽的一半
        transitionBand：滤波器过渡带带宽
        passAmp：通带增益，默认为2dB
        stopAmp：阻带衰减，默认为23dB"""
        #--------------设计滤波器参数--------------
        BPF_passFreq = [2 * (self.carrierFreq - halfBand) / self.sampleFreq, \
                        2 * (self.carrierFreq + halfBand) / self.sampleFreq]  # 滤除噪声的带通滤波器pass
        BPF_stopFreq = [2 * (self.carrierFreq - halfBand - transitionBand) / self.sampleFreq, \
                        2 * (self.carrierFreq + halfBand + transitionBand) / self.sampleFreq]  # 滤除噪声的带通滤波器stop
        maxRipple = passAmp  # 不知道是啥
        minAtten = stopAmp  # 不知道是啥

        #--------------设计滤波器阶数--------------
        BPF_order, BPF_omega_n = scisignal.ellipord(BPF_passFreq, BPF_stopFreq, passAmp, stopAmp, False)

        #--------------生成滤波器传输函数--------------
        BPF_b, BPF_a = scisignal.ellip(BPF_order, maxRipple, minAtten, BPF_omega_n, "bandpass", False, "ba")

        #--------------滤波--------------
        self.signalFilter = scisignal.filtfilt(BPF_b, BPF_a, self.signal)

        return self.signalFilter

DigiModGen=DigitalModulationGenerator

def quickMod(modType,modStep,bitsLen=1200,bitsData=None,bytesRate=100,
             carrierFreq = 1000,sampleFreq = 6000,
             seperate = 100):
    """快速生成调制信号，除调制方式外，其他参数全部默认"""
    if type(bitsLen)==int and bitsLen>0:
        bitsData=numpy.random.randint(0, 2, bitsLen)
    else:
        if bitsLen==None:
            if type(bitsData)==numpy.ndarray:
                bitsData=bitsData.copy()
            else:
                raise ValueError("bitsLen为None时必须有bitsData输入，且其为ndarray类型")
        else:
            raise ValueError("bitsLen必须为int类型，且大于0")
    DMG=DigiModGen(bitsData,bytesRate,carrierFreq,sampleFreq)
    DMG.setFSKSeperateFreq(seperate)
    constel,signal,timeSeries=DMG.get_Constel_Signal(modType,modStep)
    return constel,signal,timeSeries

if __name__=="__main__":
    import matplotlib.pyplot as plt
    import matplotlib.font_manager
    """测试用主函数"""
    bitsData = numpy.random.randint(0, 2, 16)  # 基带码元
    # bitsData=numpy.array([0,0,1,1,0,1,1,0])  # 基带码元
    bytesRate=2      # 码元速率
    carrierFreq=20    # 载波频率
    sampleFreq=60*carrierFreq   # 采样速率，应取得比较高
    ask_order=4  # ask阶数
    psk_order=4  # psk阶数
    qam_order=16  # qam阶数

    fsk_order=2  # fsk阶数
    seperate=8  # fsk两频率相较载波的偏差
    fsk_phase_continue=True

    DMG=DigitalModulationGenerator(bitsData,bytesRate,carrierFreq,sampleFreq)  # 实例化类

    askSignal,askTimeSeries=DMG.getSignal_real("ask", ask_order)  # ask信号
    pskSignal,pskTimeSeries=DMG.getSignal_real("psk", psk_order)  # psk信号
    qamSignal,qamTimeSeries=DMG.getSignal_real("qam", qam_order)  # qam信号
    fskSignal,fskTimeSeries=DMG.getSignal_real("fsk", fsk_order, True, seperate, fsk_phase_continue)  # fsk信号

    zhFont=matplotlib.font_manager.FontProperties(fname='C:\Windows\Fonts\simsun.ttc')
    fig1=plt.figure(1)
    fig1.suptitle("ASK、PSK、FSK、QAM调制信号的时域和频域波形", fontproperties = zhFont,fontsize=20)

    subplotNum=1
    plotRow=5
    plotCol=2

# -------------------------下面为绘图部分-------------------------
# --------------比特数据--------------
    ax0=fig1.add_subplot(plotRow,1,subplotNum)
    subplotNum += 2
    ax0.set_title("比特数据 f_b="+str(DMG.bytesRate)+"Hz", fontproperties = zhFont,fontsize=18)
    ax0.plot(DMG.timeSeries[:DMG.timeN],DMG.bitsData.repeat(DMG.NperBytes))
# --------------ASK--------------
    ax1=fig1.add_subplot(plotRow,plotCol,subplotNum)
    subplotNum+=1
    ax1.set_title(str(ask_order)+"ASK调制-时域", fontproperties = zhFont,fontsize=18)
    ax1.plot(askTimeSeries,askSignal)
    ax2=fig1.add_subplot(plotRow,plotCol,subplotNum)
    subplotNum += 1
    ax2.set_title(str(ask_order)+"ASK调制-FFT", fontproperties = zhFont,fontsize=18)
    ax2.plot(numpy.abs(numpy.fft.fft(askSignal)[:round(askSignal.size/2)]))
# --------------PSK--------------
    ax3=fig1.add_subplot(plotRow,plotCol,subplotNum)
    subplotNum+=1
    ax3.set_title(str(psk_order)+"PSK调制-时域", fontproperties = zhFont,fontsize=18)
    ax3.plot(pskTimeSeries,pskSignal)
    ax4=fig1.add_subplot(plotRow,plotCol,subplotNum)
    subplotNum += 1
    ax4.set_title(str(psk_order)+"PSK调制-FFT", fontproperties = zhFont,fontsize=18)
    ax4.plot(numpy.abs(numpy.fft.fft(pskSignal)[:round(pskSignal.size/2)]))
# --------------QAM--------------
    ax5=fig1.add_subplot(plotRow,plotCol,subplotNum)
    subplotNum += 1
    ax5.set_title(str(qam_order)+"QAM调制-时域", fontproperties = zhFont,fontsize=18)
    ax5.plot(qamTimeSeries,qamSignal)
    ax6=fig1.add_subplot(plotRow,plotCol,subplotNum)
    subplotNum += 1
    ax6.set_title(str(qam_order)+"QAM调制-FFT", fontproperties = zhFont,fontsize=18)
    ax6.plot(numpy.abs(numpy.fft.fft(qamSignal)[:round(qamSignal.size/2)]))
# --------------FSK--------------
    ax7=fig1.add_subplot(plotRow,plotCol,subplotNum)
    subplotNum += 1
    ax7.set_title(str(fsk_order)+"FSK调制-时域", fontproperties = zhFont,fontsize=18)
    ax7.plot(fskTimeSeries,fskSignal)
    ax8=fig1.add_subplot(plotRow,plotCol,subplotNum)
    subplotNum += 1
    ax8.set_title(str(fsk_order)+"FSK调制-FFT", fontproperties = zhFont,fontsize=18)
    ax8.plot(numpy.abs(numpy.fft.fft(fskSignal)[:round(fskSignal.size/2)]))


    fig1.show()